import tushare as ts
import pandas as pd
import matplotlib.pyplot as plt

# 初始化pro接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
#获取指定股票在指定日期范围内的日线数据
def get_stock_data(ts_code, start_date, end_date):
    df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date,
                   fields="ts_code,trade_date,open,high,low,close,pre_close,change,pct_chg,vol,amount")
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df = df.sort_values(by='trade_date')
    return df
#生成交易信号
def generate_trading_signals(df):
    df['signal'] = 0
    # 金叉信号（短期均线突破长期均线）
    short_window = 5
    long_window = 20
    df['short_mavg'] = df['close'].rolling(window=short_window).mean()
    df['long_mavg'] = df['close'].rolling(window=long_window).mean()
    df['signal'] = (df['short_mavg'] > df['long_mavg']).astype(int)
    df['position'] = df['signal'].diff()
    return df
#选股与排序
def select_and_sort_stocks(dfs):
    sorted_dfs = sorted(dfs, key=lambda x: x['pct_chg'].mean(), reverse=True)
    return sorted_dfs
def stop_loss_and_profit(df, stop_loss_ratio=0.1, take_profit_ratio=0.2):  #止损与止盈
    df['stop_loss'] = 0
    df['take_profit'] = 0
    position = 0
    buy_price = 0
    for i in range(len(df)):
        if df['position'].iloc[i] == 1:
            position = 1
            buy_price = df['close'].iloc[i]
        elif position == 1:
            if df['close'].iloc[i] <= buy_price * (1 - stop_loss_ratio):
                df.at[df.index[i], 'stop_loss'] = 1
                position = 0
            elif df['close'].iloc[i] >= buy_price * (1 + take_profit_ratio):
                df.at[df.index[i], 'take_profit'] = 1
                position = 0
    return df
#回测
def backtest(df):
    initial_capital = float(100000.0)
    positions = pd.DataFrame(index=df.index).fillna(0.0)
    positions['stock'] = df['signal']
    portfolio = positions.multiply(df['close'], axis=0)
    pos_diff = positions.diff()
    portfolio['holdings'] = (positions.multiply(df['close'], axis=0)).sum(axis=1)
    portfolio['cash'] = initial_capital - (pos_diff.multiply(df['close'], axis=0)).sum(axis=1).cumsum()
    portfolio['total'] = portfolio['cash'] + portfolio['holdings']
    portfolio['returns'] = portfolio['total'].pct_change()
    return portfolio
#可视化
def visualize(df, portfolio):
    fig, axs = plt.subplots(2, 1, figsize=(12, 8))
    axs[0].plot(df['trade_date'], df['close'], label='Close Price')
    axs[0].plot(df['trade_date'], df['short_mavg'], label='Short MA')
    axs[0].plot(df['trade_date'], df['long_mavg'], label='Long MA')
    axs[0].plot(df[df['position'] == 1]['trade_date'], df[df['position'] == 1]['close'], '^', markersize=10, color='g',
                label='Buy Signal')
    axs[0].plot(df[df['position'] == -1]['trade_date'], df[df['position'] == -1]['close'], 'v', markersize=10,
                color='r',
                label='Sell Signal')
    axs[0].set_title('Stock Price with Trading Signals')
    axs[0].set_xlabel('Date')
    axs[0].set_ylabel('Price')
    axs[0].legend()
    axs[1].plot(portfolio.index, portfolio['total'], label='Portfolio Value')
    axs[1].set_title('Portfolio Value over Time')
    axs[1].set_xlabel('Date')
    axs[1].set_ylabel('Value')
    axs[1].legend()
    plt.tight_layout()
    plt.show()
if __name__ == "__main__":
    ts_code = '000007.SZ'
    start_date = '20240401'
    end_date = '20250401'
    df = get_stock_data(ts_code, start_date, end_date)
    df = generate_trading_signals(df)
    df = stop_loss_and_profit(df)
    portfolio = backtest(df)
    visualize(df, portfolio)
